import time
from typing import List
import numpy as np
from ultralytics import YOLO
import cv2

from interfaces.detector_interface import IDetector
from models.data_models import DetectionResult, BoundingBox

class YOLOv8Detector(IDetector):
    """YOLOv8检测器实现"""
    
    def __init__(self):
        self.model = None
        self.confidence_threshold = 0.5
        self.iou_threshold = 0.45  # NMS IoU阈值
        self.max_detections = 300  # 最大检测数量
        self.agnostic_nms = False  # 类别无关NMS
        self.is_loaded = False
    
    def load_model(self, weights_path: str) -> None:
        """加载YOLOv8模型"""
        try:
            self.model = YOLO(weights_path)
            self.is_loaded = True
            print(f"模型加载成功: {weights_path}")
        except Exception as e:
            print(f"模型加载失败: {e}")
            self.is_loaded = False
            raise
    
    def detect(self, image: np.ndarray, image_path: str = "") -> DetectionResult:
        """检测图像中的天牛"""
        if not self.is_loaded:
            raise RuntimeError("模型未加载，请先调用load_model()")
        
        start_time = time.time()
        
        # 使用YOLOv8进行推理，添加更多参数来改善密集遮挡检测
        results = self.model(
            image, 
            conf=self.confidence_threshold,
            iou=self.iou_threshold,
            max_det=self.max_detections,
            agnostic_nms=self.agnostic_nms,
            verbose=False
        )
        
        boxes = []
        for result in results:
            if result.boxes is not None:
                for box in result.boxes:
                    # 提取边界框坐标和置信度
                    x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
                    confidence = box.conf[0].cpu().numpy()
                    class_id = int(box.cls[0].cpu().numpy()) if box.cls is not None else 0
                    
                    bbox = BoundingBox(
                        x1=float(x1),
                        y1=float(y1),
                        x2=float(x2),
                        y2=float(y2),
                        confidence=float(confidence),
                        class_id=class_id
                    )
                    boxes.append(bbox)
        
        processing_time = time.time() - start_time
        
        return DetectionResult(
            image_path=image_path,
            boxes=boxes,
            count=len(boxes),
            processing_time=processing_time,
            confidence_threshold=self.confidence_threshold
        )
    
    def set_confidence_threshold(self, threshold: float) -> None:
        """设置置信度阈值"""
        if 0.0 <= threshold <= 1.0:
            self.confidence_threshold = threshold
        else:
            raise ValueError("置信度阈值必须在0.0到1.0之间")
    
    def set_iou_threshold(self, threshold: float) -> None:
        """设置IoU阈值（用于NMS）"""
        if 0.0 <= threshold <= 1.0:
            self.iou_threshold = threshold
        else:
            raise ValueError("IoU阈值必须在0.0到1.0之间")
    
    def set_max_detections(self, max_det: int) -> None:
        """设置最大检测数量"""
        if max_det > 0:
            self.max_detections = max_det
        else:
            raise ValueError("最大检测数量必须大于0")
    
    def set_agnostic_nms(self, agnostic: bool) -> None:
        """设置是否使用类别无关的NMS"""
        self.agnostic_nms = agnostic
    
    def get_model_info(self) -> dict:
        """获取模型信息"""
        if not self.is_loaded:
            return {"loaded": False}
        
        return {
            "loaded": True,
            "confidence_threshold": self.confidence_threshold,
            "iou_threshold": self.iou_threshold,
            "max_detections": self.max_detections,
            "agnostic_nms": self.agnostic_nms,
            "model_type": "YOLOv8"
        }
